# Libraries Used
import os
import sys

import numpy as np # Numerical computing
import pandas as pd # Data manipulation

import seaborn as sns # R-like high-level plots
import matplotlib.pyplot as plt # Python's base plotting 

# home path
sys.path.append('../')
from src import utils, plots

from ProteoForge import annotate
from ProteoForge import proteoform_classifier as pfc

# Initialize the timer
startTime = utils.getTime()

## Configure Plotting Settings
def_colors = [
    "#139593", "#fca311", "#e54f2a",
    "#c3c3c3", "#555555",
    "#690000", "#5f4a00", "#004549"
]

condition_colors = {
    'True Hypoxia (1% O2) - 48 Hr': '#e9c46a',
    'True Hypoxia (1% O2) - 72 Hr': '#f4a261',
    'Normoxia (21% O2) - 48 Hr': '#ade8f4',
    'Normoxia (21% O2) - 72 Hr': '#0077b6',
    # Won't be used since the manuscript didn't cover these conditions
    # 'Chemical Hypoxia (CoCl2) - 48 Hr': '#c77dff',
    # 'Chemical Hypoxia (CoCl2) - 72 Hr': '#8338ec',
    # 'Oxidative Stress (H2O2) - 48 Hr': '#ff006e',
    # 'Oxidative Stress (H2O2) - 72 Hr': '#fb5607',
}

# Set seaborn style
sns.set_theme(
    style="white",
    context="paper",
    palette=def_colors,
    font_scale=1,
    rc={
        "figure.figsize": (6, 4),
        "font.family": "sans-serif",
        "font.sans-serif": ["Arial", "Ubuntu Mono"],
    }
)

# Figure Saving Settings
figure_formats = ["png", "pdf"]
save_to_folder = True
transparent_bg = True
figure_dpi = 300

## Configure dataframe displaying
pd.set_option('display.float_format', lambda x: '%.4f' % x)
pd.set_option('display.max_rows', 50)
pd.set_option('display.max_columns', 50)
pd.set_option('display.width', 500)  # Set a wider display width


## Printing Settings
verbose = True

plots.color_palette( def_colors, save=False )
plots.color_palette( condition_colors, save=False, name='condition', size=2)

home_path = './'
data_name = "hypoxia"
notebook_name = "03-ResultsAnalysis"
data_path = f"{home_path}data/"
input_path = f"{data_path}cleaned/{data_name}/"
output_path = f"{data_path}results/{data_name}/"
figure_path = f"{home_path}figures/{data_name}/{notebook_name}/"

# Create the output folder if it does not exist
if not os.path.exists(output_path):
    os.makedirs(output_path)

# Create figure folder structure, if needed
if save_to_folder:
    for i in figure_formats:
        cur_folder = figure_path + i + "/"
        if not os.path.exists(cur_folder):
            os.makedirs(cur_folder)

uniprot_data = pd.read_feather(f"{output_path}uniprot_data_Comb.feather")
uniprot_data = annotate.add_feature_metadata(uniprot_data, annotate.feature_categories, annotate.relevant_features)
print(f"UniProt data loaded with shape: {uniprot_data.shape}")
uniprot_data.info()
uniprot_data

meta = pd.read_csv(f"{input_path}metadata.csv")
# Build a dictionary to match condition_colors keys: 'Condition - Duration Hr'
condition_to_samples = {}
for condition in meta['Condition'].unique():
    for duration in sorted(meta.loc[meta['Condition'] == condition, 'Duration'].unique()):
        key = f"{condition} - {duration} Hr"
        samples = meta.loc[(meta['Condition'] == condition) & (meta['Duration'] == duration), 'Colname'].tolist()
        condition_to_samples[key] = samples

# samples_to_condition mapping
samples_to_condition = {sample: condition for condition, samples in condition_to_samples.items() for sample in samples}
# samples_to_colors mapping
samples_to_colors = {sample: condition_colors[condition] for sample, condition in samples_to_condition.items()}
meta.head(5)

pThr = 10**-3

test_data = pd.read_feather(f"{output_path}test_data_Comb.feather")
print(test_data.shape)
print()


summary_data = pd.read_feather(f"{output_path}summary_data_Comb.feather")
print(summary_data.shape)
print()

info_data = pd.read_feather(f"{input_path}info_data.feather")
print(info_data.shape)
print()

summary_data['trace'] = info_data['trace']

protein_data = pd.read_feather(f"{input_path}protein_data.feather")
print(protein_data.shape)
print()


## Adding Hypoxia Regulation compared to Normoxia 

df =  test_data.copy()

# --- 2. Define Control and Treatment Conditions ---
control_condition = 'Normoxia (21% O2) - 72 Hr'
# Sorting ensures the regulation string (e.g., 'Up_Down') is always in a consistent order.
treatment_conditions = sorted([cond for cond in df['Condition'].unique() if cond != control_condition])

print(f"Control: {control_condition}")
print(f"Treatments (in order): {treatment_conditions}")

# --- 3. Calculate Mean Intensity for Each Peptide in Each Condition ---
# This aggregates the technical replicates for each condition.
# The mean of AdjIntensity for a treatment is its log2FC vs. control.
# Using .agg is slightly more explicit and can be faster than .mean()
mean_intensities = df.groupby(['Protein', 'PeptideID', 'Condition'], as_index=False).agg(log2FC=('AdjIntensity', 'mean'))


# --- 4. Create the Regulation Table ---
# We can filter out the control condition, as its regulation is baseline (0).
regulation_df = mean_intensities[mean_intensities['Condition'] != control_condition].copy()

# Calculate the linear Fold Change from the log2FC.
# FoldChange = 2^(log2FC)
regulation_df['FoldChange'] = 2**regulation_df['log2FC']


# --- 5. Categorize Regulation Tendencies (Optimized) ---
# To categorize tendencies, we first pivot the table to have conditions as columns.
# This allows us to see the behavior of one peptide across all conditions in a single row.
pivot_reg = regulation_df.pivot_table(
    index=['Protein', 'PeptideID'],
    columns='Condition',
    values='log2FC'
).fillna(0) # Fill peptides not found in a condition with 0 log2FC (no change)

# This vectorized approach is much faster than using .apply()
# It builds the regulation string (e.g., 'Up_Down') column by column.
regulation_strings = None
for i, cond in enumerate(treatment_conditions):
    # Use np.select for fast conditional categorization based on direction
    categories = np.select(
        [pivot_reg[cond] > 0, pivot_reg[cond] < 0],
        ['Up', 'Down'],
        default='NC' # No Change
    )
    
    # Convert to a pandas Series to align with the pivot_reg index
    categories_series = pd.Series(categories, index=pivot_reg.index)

    if i == 0:
        regulation_strings = categories_series
    else:
        # Append the category for the next condition
        regulation_strings += '_' + categories_series

pivot_reg['Regulation'] = regulation_strings


# Now, merge the 'Regulation' column back into our long-format table
# This adds the summary string to each row corresponding to a given peptide.
regulation_df = regulation_df.merge(
    pivot_reg[['Regulation']],
    on=['Protein', 'PeptideID'],
    how='left'
)

# Reorder columns for a clean final output
regulation_df = regulation_df[['Protein', 'PeptideID', 'Condition', 'log2FC', 'FoldChange', 'Regulation']]


print("\n--- Final Regulation Table with Tendencies ---")
print(regulation_df.head())

# You can now save this final dataframe to a new file, for example:
# regulation_df.to_csv('regulation_analysis_output.csv', index=False)

# Example of how to use the new 'Regulation' column:
print("\n--- Example: Grouping by Regulation Tendency ---")
print(regulation_df.groupby('Regulation')['PeptideID'].nunique())

plot_data = regulation_df[['Protein', 'PeptideID', 'Regulation']].drop_duplicates().copy()

# Adding the plot_data to summary_data to get the Regulation column (use hypoxia 48hr-72hr)
summary_data = summary_data.merge(
    plot_data[['Protein', 'PeptideID', 'Regulation']],
    on=['Protein', 'PeptideID'],
    how='left'
)
summary_data['Regulation'] = summary_data['Regulation'].fillna('NC_NC')  # peptides not in regulation_df are 'No Change' in both conditions
summary_data.rename(columns={'Regulation': 'Regulation_(Hypoxia 48hr-72hr)'}, inplace=True)

fig, ax = plt.subplots(figsize=(8, 6))
plot_data['Regulation'].value_counts().plot(kind='barh', color="#432c46", edgecolor='black', ax=ax)
# Add counts to on top of the bars for barh
for i, v in enumerate(plot_data['Regulation'].value_counts()):
    ax.text(
        v, i, str(v),
        color='black', 
        ha='left', va='center', 
        rotation=0, fontsize=15, fontweight='bold'
    )
ax.set_xlabel("Number of Peptides", fontsize=12)
ax.set_ylabel("Regulation Tendency (Hypoxia 48Hr-72Hr)", fontsize=12)
ax.set_title("Peptide Regulation Tendencies Across Conditions", fontsize=14, fontweight='bold')
ax.grid("both", linestyle="--", linewidth=0.75, alpha=0.5, color="lightgrey")
plt.tight_layout()

plots.finalize_plot( 
    fig, show=True, save=save_to_folder, 
    filename="peptide_regulation_tendencies",
    filepath=figure_path,
    formats=figure_formats,
    transparent=transparent_bg,
    dpi=figure_dpi
)

print()
summary_data['Regulation_(Hypoxia 48hr-72hr)'].value_counts()

# Create protein-level summary with dPF categories (vectorized, fast pandas)
# Uses indicator columns and a single groupby.agg to avoid per-group Python callbacks.

# local alias
df = summary_data.copy()

# Ensure numeric/bool-like columns use compact dtypes
if df['isSignificant'].dtype != 'int32' and df['isSignificant'].dtype != 'int64':
    df['isSignificant'] = df['isSignificant'].fillna(False).astype('int32')

# Temporary indicator columns (fast to sum)
df['is_canonical'] = (df['dPF'] == 0).astype('int32')
df['is_ptm'] = (df['dPF'] == -1).astype('int32')
df['is_proteoform'] = (df['dPF'] > 0).astype('int32')

# Perform one grouped aggregation (all C-level operations)
protein_dpf_analysis = (
    df.groupby('Protein', sort=False, observed=True)
      .agg(
          Gene = ('Gene', 'first'),
          Length = ('Length', 'first'),
          Coverage = ('Coverage', 'first'),
          all_pep = ('isSignificant', 'count'),
          sgnf_pep = ('isSignificant', 'sum'),
          all_clust = ('ClusterID', 'max'),
          n_forms = ('dPF', 'nunique'),
          n_canon = ('is_canonical', 'sum'),
          n_ptm = ('is_ptm', 'sum'),
          n_dPF = ('is_proteoform', 'sum'),
          max_dpf = ('dPF', 'max'),
      )
)

# Convert max_dpf -> max_proteoform_id: only positive dPFs indicate proteoform ids
protein_dpf_analysis['n_forms'] = (
    protein_dpf_analysis['max_dpf']
      .fillna(0)
      .where(protein_dpf_analysis['max_dpf'] > 0, 0)
      .astype('int32')
)
protein_dpf_analysis.drop(columns=['max_dpf'], inplace=True)

# If other code expects a flattened/explicit column order, optionally reorder here.
# Cleanup temporary indicator columns to free memory
try:
    df.drop(columns=['is_canonical', 'is_ptm', 'is_proteoform'], inplace=True)
except Exception:
    pass

# Add population rates for peptides, canonical, PTM, and proteoforms
protein_dpf_analysis['sgnf_pct'] = (
    protein_dpf_analysis['sgnf_pep'] / protein_dpf_analysis['all_pep']
).fillna(0).astype('float32') * 100.0 
protein_dpf_analysis['canon_pct'] = (
    protein_dpf_analysis['n_canon'] / protein_dpf_analysis['all_pep']
).fillna(0).astype('float32') * 100.0 
protein_dpf_analysis['ptm_pct'] = (
    protein_dpf_analysis['n_ptm'] / protein_dpf_analysis['all_pep']
).fillna(0).astype('float32') * 100.0 
protein_dpf_analysis['dPF_pct'] = (
    protein_dpf_analysis['n_dPF'] / protein_dpf_analysis['all_pep']
).fillna(0).astype('float32') * 100.0 
protein_dpf_analysis['eff_pct'] = (
    protein_dpf_analysis['n_forms'] / protein_dpf_analysis['all_clust']
).fillna(0).astype('float32') * 100.0

# Categorize coverage
protein_dpf_analysis['cov_cat'] = pd.cut(
    protein_dpf_analysis['Coverage'],
    bins=[-1, 25, 50, 75, 100],
    labels=['Low', 'Medium', 'High', 'V.High'],
    include_lowest=True
)
# Categorize available significant peptides per protein
protein_dpf_analysis['sgnf_cat'] = pd.cut(
    protein_dpf_analysis['sgnf_pep'],
    bins=[-1, 0, 1, np.inf],
    labels=['None', 'Single', 'Multiple'],
    right=True
)
# Categorize protein length
protein_dpf_analysis['len_cat'] = pd.cut(
    protein_dpf_analysis['Length'],
    bins=[-1, 150, 250, 750, 1000, 2000, np.inf],
    labels=['V.Short', 'Short', 'Medium', 'Long', 'V.Long', 'Huge'],
    right=True
)


# Categorize dPFs and PTMs containment per protein
labels = ['No Forms', 'PTM Only', 'dPF Only', 'Both']

conds = [
    (protein_dpf_analysis['n_forms'] == 0) & (protein_dpf_analysis['n_ptm'] == 0),
    (protein_dpf_analysis['n_forms'] == 0) & (protein_dpf_analysis['n_ptm'] > 0),
    (protein_dpf_analysis['n_forms'] > 0) & (protein_dpf_analysis['n_ptm'] == 0),
    (protein_dpf_analysis['n_forms'] > 0) & (protein_dpf_analysis['n_ptm'] > 0),
]

protein_dpf_analysis['dpf_cat'] = pd.Categorical(
    np.select(conds, labels, default='No Forms'),
    categories=labels
)

# protein_dpf_analysis is now ready with ptm_count, proteoform_count, max_proteoform_id, etc.
protein_dpf_analysis = protein_dpf_analysis.reset_index()
protein_dpf_analysis

# Heatmap of average sgnf_pct by coverage category vs number of proteoforms (n_forms)
heatmap_data = protein_dpf_analysis.pivot_table(
    values='sgnf_pct',
    index='cov_cat',
    columns='n_forms',
    aggfunc='mean',
    observed=True,
)

# Ensure sensible ordering of coverage categories if categorical
if isinstance(protein_dpf_analysis['cov_cat'].dtype, pd.CategoricalDtype):
    heatmap_data = heatmap_data.reindex(protein_dpf_analysis['cov_cat'].cat.categories)

# Sort columns (n_forms) ascending
heatmap_data = heatmap_data.reindex(sorted(heatmap_data.columns), axis=1)

plt.figure(figsize=(10, 6))
sns.heatmap(
    heatmap_data,
    annot=True,
    fmt=".1f",
    cmap="YlGnBu",
    cbar_kws={'label': 'Average sgnf_pct (%)'},
    linewidths=0.5,
    linecolor='white',
    mask=heatmap_data.isna()
)
plt.title('Average Significance Percentage (sgnf_pct) by Coverage Category vs Number of Proteoforms')
plt.ylabel('Coverage Category')
plt.xlabel('Number of Proteoforms (n_forms)')
plt.tight_layout()
plots.finalize_plot( 
    fig, show=True, save=save_to_folder, 
    filename='heatmap_sgnf_pct_cov_cat_vs_n_forms',
    filepath=figure_path,
    formats=figure_formats,
    transparent=transparent_bg,
    dpi=figure_dpi
)

# Plot Composition of Proteins by dPF Categories
# This plot shows the distribution of canonical, PTM, and proteoform peptides across proteins
# with a smoothed coverage trend line.
custom_colors = {
    'Canonical': "#e5e5e5",      # Gray - baseline
    'PTM': "#a7c957",            # Red - modifications  
    'Proteoform': "#bc4749",     # Blue - true proteoforms
}

plot_data = protein_dpf_analysis.copy()
proteins_sorted = plot_data.sort_values('Coverage').reset_index(drop=True)
n_proteins = len(proteins_sorted)

# If many proteins, aggregate into bins (keeps trend readable)
# tune: lower -> more smoothing, higher -> more detail
max_bins = 100 
if n_proteins > max_bins:
    # group into evenly spaced bins
    proteins_sorted['bin'] = pd.cut(np.arange(n_proteins), bins=max_bins, labels=False)
    agg = proteins_sorted.groupby('bin').agg({
        'canon_pct':'mean',
        'ptm_pct':'mean',
        'dPF_pct':'mean',
        'Coverage':'mean'
    }).reset_index(drop=True)
    x = np.arange(len(agg))
    canonical = agg['canon_pct'].values
    ptm = agg['ptm_pct'].values
    proteoform = agg['dPF_pct'].values
    coverage = agg['Coverage'].values
    xtick_locs = np.linspace(0, len(agg)-1, 5).astype(int)
    xtick_labels = [f"{int(i/len(agg)*100)}%" for i in xtick_locs]  # percent position
else:
    x = np.arange(n_proteins)
    canonical = proteins_sorted['canon_pct'].values
    ptm = proteins_sorted['ptm_pct'].values
    proteoform = proteins_sorted['dPF_pct'].values
    coverage = proteins_sorted['Coverage'].values
    xtick_locs = np.linspace(0, n_proteins-1, 5).astype(int)
    xtick_labels = [f"{int(i/ n_proteins * 100)}%" for i in xtick_locs]

# Compute overall means for legend
mean_c = canonical.mean()
mean_ptm = ptm.mean()
mean_pf = proteoform.mean()

fig, ax = plt.subplots(figsize=(12, 5))
# stackplot handles the stacking cleanly
ax.stackplot(
    x,
    canonical, ptm, proteoform,
    labels=[f"Canonical ({mean_c:.1f}%)",
            f"PTMs ({mean_ptm:.1f}%)",
            f"Proteoforms ({mean_pf:.1f}%)"],
    colors=[custom_colors['Canonical'], custom_colors['PTM'], custom_colors['Proteoform']],
    alpha=0.85
)

# Coverage trend on twin axis (smoothed)
ax2 = ax.twinx()
window = max(3, int(len(coverage) * 0.02))  # ~2% window, min 3
coverage_smooth = pd.Series(coverage).rolling(window=window, center=True, min_periods=1).mean().values
ax2.plot(x, coverage_smooth, color='k', linewidth=1.6, alpha=0.9, label='Coverage (smoothed)')
ax2.set_ylim(0, 100)
ax.set_xlim(x.min(), x.max())

# Tidy up ticks/labels
ax.set_xlabel(f'{plot_data.shape[0]} Proteins (sorted by coverage)', fontsize=12, fontweight='bold')
ax.set_ylabel('Peptide composition (%)', fontsize=12, fontweight='bold')
ax2.set_ylabel('Coverage (%)', fontsize=12, fontweight='bold')
ax.set_xticks(xtick_locs)
ax.set_xticklabels(xtick_labels)
ax.set_ylim(0, max( (canonical + ptm + proteoform).max(), 100))  # sensible y-limits

# Legend and annotation
leg = ax.legend(
    loc='lower right', fontsize=10, frameon=False, bbox_to_anchor=(1.0, 0.1), 
    title='Peptide Composition', title_fontsize='12', ncol=1
)
ax2.legend(loc='lower right', fontsize=10, frameon=False, bbox_to_anchor=(1.0, 0))

# Add light grid and reduce spines clutter
ax.grid(axis='y', linestyle='--', linewidth=0.6, alpha=0.75, color='gray')

plt.tight_layout()
plots.finalize_plot( 
    fig, show=True, save=save_to_folder, 
    filename='peptideComposition_proteinCoverage',
    filepath=figure_path,
    formats=figure_formats,
    transparent=transparent_bg,
    dpi=figure_dpi
)

# Plot the Efficiency of Proteoform Discovery
# X: Number of Peptides per Protein
# Y: Number of Proteoforms Discovered (non-)
# Size: Protein Length (bubble size)
# Color: Coverage
# Alpha: Coverage

from matplotlib.lines import Line2D

plot_data = protein_dpf_analysis.copy()
# plot_data['discovery_efficiency'] = plot_data['n_dPF'] / plot_data['all_pep'] # dPF_pct
# plot_data['significance_efficiency'] = plot_data['n_sgnf'] / plot_data['all_pep'] # sgnf_pct
plot_data['complexity_score'] = (plot_data['n_forms'] * plot_data['all_clust']) / (plot_data['Length'] / 100)
plot_data['complexity_score'] = plot_data['complexity_score'].fillna(0)

# Bubble size scaling (marker area)
length = plot_data['Length'].clip(lower=1)
min_area, max_area = 40, 1600
areas = np.interp(length, (length.min(), length.max()), (min_area, max_area))

# Coverage -> color and per-point alpha
coverage = plot_data['Coverage'].fillna(0).astype(float)
# cov_min, cov_max = coverage.min(), coverage.max()
cov_min, cov_max = 0, 100  # assuming coverage is in percentage
if cov_max == cov_min:
    cov_max = cov_min + 1e-6  # avoid zero-range

# map coverage to alpha range (low coverage -> lower alpha)
alpha_min, alpha_max = 0.30, 0.95
alphas = np.interp(coverage, (cov_min, cov_max), (alpha_min, alpha_max))

# build RGBA colors from a cmap, then set per-point alpha
cmap = plt.cm.RdYlBu_r
norm = plt.Normalize(vmin=cov_min, vmax=cov_max)
base_colors = cmap(norm(coverage))  # Nx4 array
base_colors[:, 3] = alphas         # set alpha channel per-point

# Plot
x = plot_data['all_pep'].values
y = plot_data['n_forms'].values

fig, ax = plt.subplots(figsize=(10, 6))

sc = ax.scatter(
    x, y,
    s=areas,
    c=base_colors,
    edgecolors='white',
    linewidth=0.8,
    zorder=2
)

# Colorbar for Coverage (note: alpha not reflected in bar)
mappable = plt.cm.ScalarMappable(norm=norm, cmap=cmap)
mappable.set_array(coverage)
cbar = fig.colorbar(mappable, ax=ax, pad=0.02, shrink=0.9)
cbar.set_label('Coverage (%)', fontsize=11, fontweight='bold')

# Size legend (approximate protein lengths)
lens = np.percentile(length, [10, 50, 90])
proxy_areas = np.interp(lens, (length.min(), length.max()), (min_area, max_area))
legend_proxies = [
    Line2D([0], [0], marker='o', color='w', label=f'Len ≈ {int(lens[0])}', markerfacecolor='gray', markersize=np.sqrt(proxy_areas[0])),
    Line2D([0], [0], marker='o', color='w', label=f'Len ≈ {int(lens[1])}', markerfacecolor='gray', markersize=np.sqrt(proxy_areas[1])),
    Line2D([0], [0], marker='o', color='w', label=f'Len ≈ {int(lens[2])}', markerfacecolor='gray', markersize=np.sqrt(proxy_areas[2]))
]

# Legends: exceptional + size proxies
main_handles, main_labels = ax.get_legend_handles_labels()
# ax.legend(loc='upper left', fontsize=10)
ax2 = ax.legend(handles=legend_proxies, title='Protein length (approx)', loc='upper right', fontsize=9, title_fontsize=10, frameon=False)
ax.add_artist(ax2)

# Labels and title
ax.set_xlabel('Number of Peptides per Protein', fontsize=12, fontweight='bold')
ax.set_ylabel('Number of Proteoforms Discovered', fontsize=12, fontweight='bold')
ax.set_title(
    'Proteoform Discovery Landscape\n(Size ~ Protein length, Color = Coverage; alpha ~ Coverage)',
    fontsize=14, fontweight='bold'
)
ax.grid(True, which="both", linestyle="--", linewidth=0.5, alpha=0.7)

plt.tight_layout()
plots.finalize_plot( 
    fig, show=True, save=save_to_folder, 
    filename='proteoform_discovery_efficiency',
    filepath=figure_path,
    formats=figure_formats,
    transparent=transparent_bg,
    dpi=figure_dpi
)

# Revised layout: stacked bar (top) + percentage lines + efficiency (bottom) with shared x-axis
import matplotlib.patches as mpatches

plot_data = protein_dpf_analysis.copy()

# 1. Prepare data (reuse same ordering)
composition_data = pd.crosstab(plot_data['all_clust'], plot_data['dpf_cat'])
category_order = ['dPF Only', 'PTM Only', 'No Forms', 'Both']
composition_data = composition_data.reindex(columns=category_order, fill_value=0).sort_index()
totals = composition_data.sum(axis=1).astype(int)
x_coords = np.arange(len(composition_data))

# colors in same order as columns
col_color_map = {"No Forms":"#e5e5e5","PTM Only":"#a7c957","dPF Only":"#bc4749","Both":"#967aa1"}
bar_colors = [col_color_map[c] for c in composition_data.columns]

# 2. Create stacked bar (top) and line/CI plot (bottom) sharing X
fig, (ax_bar, ax_line) = plt.subplots(2, 1, figsize=(12, 8), sharex=True, gridspec_kw={'height_ratios':[1, 1]})

# Top: stacked bars
composition_data.plot(
    kind='bar',
    stacked=True,
    ax=ax_bar,
    color=bar_colors,
    width=0.8,
    legend=False
)
ax_bar.set_ylabel('Number of Proteins', fontsize=12)
ax_bar.set_title('Composition of Proteoform Discoveries by Cluster Count', fontsize=16, pad=12)
ax_bar.grid(True, which="both", linestyle="--", linewidth=0.5, alpha=0.7)
ax_bar.set_ylim(0, int(totals.max() * 1.12))

# totals above bars
y_offset = max(2, totals.max() * 0.02)
for i, total in enumerate(totals.values):
    if total > 0:
        ax_bar.text(i, total + y_offset, f"{total:,}", ha='center', va='bottom', fontsize=10, fontweight='bold')

# tidy x tick labels on the bottom axis
ax_line.set_xlabel('Number of Clusters (all_clust)', fontsize=12)
ax_line.set_xticks(x_coords)
ax_line.set_xticklabels([str(x) for x in composition_data.index], rotation=0)

# 3. Bottom: percentage lines per category with 95% CI (binomial) + efficiency mean with CI (t-approx)
# Percentage per category (p) and 95% CI using binomial approximation: CI = 1.96 * sqrt(p*(1-p)/N)
for col, color in zip(composition_data.columns, bar_colors):
    counts = composition_data[col].values.astype(float)
    N = totals.values.astype(float)
    # avoid division by zero
    p = np.divide(counts, N, out=np.zeros_like(counts, dtype=float), where=N!=0)
    sem_binom = np.sqrt(np.divide(p * (1 - p), np.where(N == 0, 1, N)))
    ci95 = sem_binom * 1.96
    ax_line.errorbar(
        x_coords, p * 100,
        yerr=ci95 * 100,
        marker='o', linewidth=4, color=color, capsize=8, label=f"{col} (%)", 
        linestyle='-', markersize=8, alpha=0.8, elinewidth=2,
        markeredgecolor='black', capthick=2,
    )

# Efficiency stats by cluster: mean and 95% CI from group std / sqrt(n)
eff_stats = (
    plot_data.reset_index()
             .groupby('all_clust')['eff_pct']
             .agg(['mean', 'std', 'count'])
             .reindex(composition_data.index)
)
eff_mean = eff_stats['mean'].fillna(0)
eff_std = eff_stats['std'].fillna(0)
eff_count = eff_stats['count'].fillna(0).replace(0, np.nan)  # avoid div-by-zero
eff_sem = eff_std / np.sqrt(eff_count)
eff_ci95 = (eff_sem * 1.96).fillna(0)

# Plot efficiency mean and CI as emphasized black line + shaded band
ax_line.plot(x_coords, eff_mean.values, marker='s', color='k', linewidth=2.5, markersize=7, label='Avg Efficiency (eff_pct)')
ax_line.fill_between(x_coords, (eff_mean - eff_ci95).values, (eff_mean + eff_ci95).values, color='k', alpha=0.12)

# 4. Styling & legends
ax_line.set_ylabel('Percentage / Efficiency (%)', fontsize=12)
ax_line.set_ylim(0, 100)
ax_line.grid(True, which="both", linestyle="--", linewidth=0.5, alpha=0.7)

# Bar legend (from patches) and line legend
bar_handles = [mpatches.Patch(color=col_color_map[c], label=c) for c in composition_data.columns]
leg1 = ax_bar.legend(handles=bar_handles, title='Proteoform Category', bbox_to_anchor=(1.02, 1), loc='upper left', frameon=False)

leg2 = ax_line.legend(loc='upper center', bbox_to_anchor=(0.5, -0.18), ncol=5, frameon=False, title='Percent lines / Efficiency')

# keep bar legend visible
ax_bar.add_artist(leg1)

plt.tight_layout()
plots.finalize_plot( 
    fig, show=True, save=save_to_folder, 
    filename='proteoform_discovery_efficiency_by_cluster',
    filepath=figure_path,
    formats=figure_formats,
    transparent=transparent_bg,
    dpi=figure_dpi
)

# Plot donut chart for proteins with no signf, single signf, multiple signf
plot_data = protein_dpf_analysis[[ 'Gene', 'sgnf_pep' ]].copy()

# Categorize proteins
labels = ['None', 'Single', 'Multiple']
plot_data['sgnf_cat'] = pd.cut(
    plot_data['sgnf_pep'],
    bins=[-1, 0, 1, np.inf],
    labels=labels,
    right=True
)

# Ordered counts and percentages
counts = plot_data['sgnf_cat'].value_counts().reindex(labels, fill_value=0)
total = counts.sum()
percent = counts / total

# slightly separate very small slices
explode = [0.02 if p < 0.05 else 0.0 for p in percent]

# Use a list of color strings (in the same order as `labels`) instead of a dict.
# colors_list = ['#c3c3c3', '#fca311', '#e54f2a'] # Autumn palette
colors_list = ["#e5e5e5", "#669bbc", "#003049"] # Blue-Gray palette

fig, ax = plt.subplots(figsize=(8, 6))
wedges, texts = ax.pie(
    counts.values,
    labels=None,             
    startangle=140,
    colors=colors_list,
    wedgeprops=dict(width=0.40, edgecolor='w'),
    explode=explode,
    counterclock=False
)

# Add percent labels placed on wedges
bbox_props = dict(boxstyle="round,pad=0.3", fc="white", ec="none", alpha=0.8)
for w, p in zip(wedges, percent):
    angle = (w.theta2 + w.theta1) / 2.0
    x, y = w.r * 0.75 * np.cos(np.deg2rad(angle)), w.r * 0.75 * np.sin(np.deg2rad(angle))
    ax.text(x, y, f"{p:.1%}", ha='center', va='center', fontsize=15, bbox=bbox_props, fontweight='bold')

# Center annotation: total proteins
ax.text(0, 0, f"Total\n{total}", ha='center', va='center', fontsize=15, fontweight='bold')

# Legend with counts and percentages
legend_labels = [f"{lab}: {cnt:,} ({pct:.1%})" for lab, cnt, pct in zip(labels, counts, percent)]
ax.legend(
    wedges, legend_labels,
    title="Number of Significant Peptides",
    loc="upper center",
    bbox_to_anchor=(0.5, -0.12),
    ncol=len(labels),
    fontsize=11,
    title_fontsize=12,
    frameon=False, 
    handletextpad=0.5,
    columnspacing=1.5,
)

ax.axis('equal')
plt.title('Distribution of Proteins by Number of \nSignificant Peptides from Linear Model', fontsize=16, fontweight='bold')
plt.tight_layout()
plots.finalize_plot( 
    fig, show=True, save=save_to_folder, 
    filename='significantPeptides_perProtein_Donut',
    filepath=figure_path,
    formats=figure_formats,
    transparent=transparent_bg,
    dpi=figure_dpi
)

# Plot donut chart for proteins with no additional forms, only dPFs, only PTMs, with Both
plot_data = protein_dpf_analysis[[ 'Gene', 'n_forms', 'n_canon', 'n_ptm', 'n_dPF' ]].copy()

without_forms = set(plot_data[(plot_data['n_forms'] == 0)].index)
with_PTM_only = set(plot_data[(plot_data['n_forms'] == 0) & (plot_data['n_ptm'] > 0)].index)
with_dPF_only = set(plot_data[(plot_data['n_forms'] > 0) & (plot_data['n_ptm'] == 0)].index)
with_both = set(plot_data[(plot_data['n_forms'] > 0) & (plot_data['n_ptm'] > 0)].index)
without_forms = without_forms.difference(with_PTM_only)
print(f"Proteins with no additional forms: {len(without_forms):,}")
print(f"Proteins with only PTMs: {len(with_PTM_only):,}")
print(f"Proteins with only dPFs: {len(with_dPF_only):,}")
print(f"Proteins with both PTMs and dPFs: {len(with_both):,}")

# Categorize proteins
labels = ['No Forms', 'PTM Only', 'dPF Only', 'Both']
counts = [
    len(without_forms),
    len(with_PTM_only),
    len(with_dPF_only),
    len(with_both)
]
# Calculate percentages
percent = np.array(counts) / sum(counts)

# slightly separate very small slices
explode = [0.02 if p < 0.05 else 0.0 for p in percent]
# Use a list of color strings (in the same order as `labels`) instead of a dict.
colors_list = ["#e5e5e5", "#a7c957", "#bc4749", "#967aa1"]

fig, ax = plt.subplots(figsize=(8, 6))
wedges, texts = ax.pie(
    counts,
    labels=None,             
    startangle=140,
    colors=colors_list,
    wedgeprops=dict(width=0.40, edgecolor='w'),
    explode=explode,
    counterclock=False
)
# Add percent labels placed on wedges
bbox_props = dict(boxstyle="round,pad=0.3", fc="white", ec="none", alpha=0.8)
for w, p in zip(wedges, percent):
    angle = (w.theta2 + w.theta1) / 2.0
    x, y = w.r * 0.75 * np.cos(np.deg2rad(angle)), w.r * 0.75 * np.sin(np.deg2rad(angle))
    ax.text(x, y, f"{p:.1%}", ha='center', va='center', fontsize=15, bbox=bbox_props, fontweight='bold')

# Center annotation: total proteins
ax.text(0, 0, f"Total\n{sum(counts)}", ha='center', va='center', fontsize=15, fontweight='bold')

# Legend with counts and percentages
legend_labels = [f"{lab}: {cnt:,} ({pct:.1%})" for lab, cnt, pct in zip(labels, counts, percent)]
ax.legend(
    wedges, legend_labels,
    title="Number of Proteins",
    loc="upper center",
    bbox_to_anchor=(0.5, -0.12),
    ncol=len(labels),
    fontsize=11,
    title_fontsize=12,
    frameon=False, 
    handletextpad=0.5,
    columnspacing=1.5,
)
ax.axis('equal')
plt.title('Distribution of Proteins by Additional Forms and PTMs', fontsize=16, fontweight='bold')
plt.tight_layout()
plots.finalize_plot( 
    fig, show=True, save=save_to_folder, 
    filename='formCategories_perProtein_Donut',
    filepath=figure_path,
    formats=figure_formats,
    transparent=transparent_bg,
    dpi=figure_dpi
)

def create_enhanced_peptide_uniprot_map_with_miscleavage(
    summary_data: pd.DataFrame,
    uniprot_data: pd.DataFrame,
    protein_dpf_analysis: pd.DataFrame,
    overlap_threshold: float = 0.5,
    verbose: bool = False
) -> pd.DataFrame:
    """
    Create an enhanced peptide-UniProt annotation mapping with miscleavage detection.

    Parameters
    ----------
    summary_data : pd.DataFrame
        Peptide summary data with columns: Protein, PeptideID, peptide_start, peptide_end, dPF.
    uniprot_data : pd.DataFrame
        UniProt annotation data with columns: Protein, feature_category, feature, start, end, group, agent, note, score, is_relevant.
    protein_dpf_analysis : pd.DataFrame
        Protein dPF analysis with at least columns: Protein, dpf_cat.
    overlap_threshold : float, optional
        Fractional overlap threshold for miscleavage detection (default is 0.5).
    verbose : bool, optional
        If True, print detailed progress and statistics.

    Returns
    -------
    pd.DataFrame
        DataFrame with columns:
        Protein, PeptideID, pepStart, pepEnd, startpos, endpos, feature_category, feature,
        group, agent, note, score, dPF.
    """
    import time
    t0 = time.time()
    if verbose:
        print("=== Enhanced Peptide-UniProt Mapping with Miscleavage Detection ===")

    # Filter proteins of interest
    proteins = protein_dpf_analysis.loc[
        protein_dpf_analysis['dpf_cat'] != 'No Forms', 'Protein'
    ].unique()
    if verbose:
        print(f"  Proteins to check: {len(proteins):,}")

    # Prepare peptides
    peptides = (
        summary_data[summary_data['Protein'].isin(proteins)]
        .rename(columns={'peptide_start': 'pepStart', 'peptide_end': 'pepEnd'})
        [['Protein', 'PeptideID', 'pepStart', 'pepEnd', 'dPF']]
        .reset_index(drop=True)
    )
    if verbose:
        print(f"  Peptides to process: {len(peptides):,}")

    # Prepare relevant UniProt annotations
    annotations = (
        uniprot_data[
            (uniprot_data['is_relevant']) &
            (uniprot_data['Protein'].isin(proteins))
        ]
        .rename(columns={'start': 'ann_start', 'end': 'ann_end'})
        [['Protein', 'feature_category', 'feature', 'ann_start', 'ann_end', 'group', 'agent', 'note', 'score']]
        .reset_index(drop=True)
    )
    if verbose:
        print(f"  Relevant UniProt annotations: {len(annotations):,}")

    # Find peptide-annotation overlaps
    merged = peptides.merge(annotations, on='Protein', how='left')
    overlap_mask = (
        ((merged['ann_start'] >= merged['pepStart']) & (merged['ann_start'] <= merged['pepEnd'])) |
        ((merged['ann_end'] >= merged['pepStart']) & (merged['ann_end'] <= merged['pepEnd'])) |
        ((merged['ann_start'] <= merged['pepStart']) & (merged['ann_end'] >= merged['pepEnd']))
    )
    ann_overlap = merged[overlap_mask].copy()

    # Build annotation result
    result_ann = pd.DataFrame()
    if not ann_overlap.empty:
        result_ann = pd.DataFrame({
            'Protein': ann_overlap['Protein'],
            'PeptideID': ann_overlap['PeptideID'],
            'pepStart': ann_overlap['pepStart'],
            'pepEnd': ann_overlap['pepEnd'],
            'startpos': ann_overlap['ann_start'],
            'endpos': ann_overlap['ann_end'],
            'feature_category': ann_overlap['feature_category'],
            'feature': ann_overlap['feature'],
            'group': ann_overlap['group'],
            'agent': ann_overlap['agent'],
            'note': ann_overlap['note'],
            'score': ann_overlap['score'],
            'dPF': ann_overlap['dPF']
        })

    # Miscleavage/peptide-peptide overlap detection
    if verbose:
        print("  Detecting miscleavage/peptide-peptide overlaps...")

    pep2pep = peptides.merge(peptides, on='Protein', suffixes=('', '_other'))
    pep2pep = pep2pep[pep2pep['PeptideID'] != pep2pep['PeptideID_other']]
    pep2pep['pep_length'] = pep2pep['pepEnd'] - pep2pep['pepStart'] + 1
    pep2pep['overlap_start'] = np.maximum(pep2pep['pepStart'], pep2pep['pepStart_other'])
    pep2pep['overlap_end'] = np.minimum(pep2pep['pepEnd'], pep2pep['pepEnd_other'])
    actual_overlap = pep2pep[pep2pep['overlap_start'] <= pep2pep['overlap_end']].copy()
    actual_overlap['overlap_length'] = actual_overlap['overlap_end'] - actual_overlap['overlap_start'] + 1
    actual_overlap['overlap_fraction'] = actual_overlap['overlap_length'] / actual_overlap['pep_length']
    sig_overlap = actual_overlap[actual_overlap['overlap_fraction'] > overlap_threshold].copy()
    sig_overlap['overlap_type'] = np.where(sig_overlap['overlap_fraction'] > 0.8, 'Miscleavage', 'Partial_Overlap')

    result_mis = pd.DataFrame()
    if not sig_overlap.empty:
        result_mis = pd.DataFrame({
            'Protein': sig_overlap['Protein'],
            'PeptideID': sig_overlap['PeptideID'],
            'pepStart': sig_overlap['pepStart'],
            'pepEnd': sig_overlap['pepEnd'],
            'startpos': sig_overlap['overlap_start'],
            'endpos': sig_overlap['overlap_end'],
            'feature_category': 'Cleavage Analysis',
            'feature': 'MISCLEAVAGE',
            'group': sig_overlap['overlap_type'],
            'agent': 'ProteoForge_Analysis',
            'note': 'Overlap with PeptideID ' + sig_overlap['PeptideID_other'].astype(str) +
                    ' (fraction: ' + sig_overlap['overlap_fraction'].round(3).astype(str) + ')',
            'score': sig_overlap['overlap_fraction'],
            'dPF': sig_overlap['dPF']
        })

    # Peptides with no annotation or miscleavage
    all_pep_ids = set(zip(peptides['Protein'], peptides['PeptideID']))
    annotated_ids = set()
    if not result_ann.empty:
        annotated_ids.update(zip(result_ann['Protein'], result_ann['PeptideID']))
    if not result_mis.empty:
        annotated_ids.update(zip(result_mis['Protein'], result_mis['PeptideID']))
    no_annot_ids = all_pep_ids - annotated_ids

    result_none = pd.DataFrame()
    if no_annot_ids:
        no_annot_df = peptides[
            peptides.set_index(['Protein', 'PeptideID']).index.isin(no_annot_ids)
        ].copy()
        result_none = pd.DataFrame({
            'Protein': no_annot_df['Protein'],
            'PeptideID': no_annot_df['PeptideID'],
            'pepStart': no_annot_df['pepStart'],
            'pepEnd': no_annot_df['pepEnd'],
            'startpos': -1,
            'endpos': -1,
            'feature_category': '',
            'feature': 'No Annotation',
            'group': '',
            'agent': '',
            'note': '',
            'score': 0.0,
            'dPF': no_annot_df['dPF']
        })

    # Combine all results
    result = pd.concat([df for df in [result_ann, result_mis, result_none] if not df.empty], ignore_index=True)

    # Standardize dtypes
    dtype_map = {
        'Protein': 'string',
        'PeptideID': 'int32',
        'pepStart': 'int32',
        'pepEnd': 'int32',
        'startpos': 'int32',
        'endpos': 'int32',
        'feature_category': 'string',
        'feature': 'string',
        'group': 'string',
        'agent': 'string',
        'note': 'string',
        'score': 'float32',
        'dPF': 'int8'
    }
    for col, typ in dtype_map.items():
        if col in result.columns:
            result[col] = result[col].astype(typ)
    string_cols = ['feature_category', 'group', 'agent', 'note']
    result[string_cols] = result[string_cols].fillna('')

    # Verbose summary
    if verbose:
        t1 = time.time()
        miscleavage_count = (result['feature'] == 'MISCLEAVAGE').sum()
        annotation_count = (result['feature'] != 'MISCLEAVAGE').sum()
        no_annotation_count = (result['feature'] == 'No Annotation').sum()
        print(f"\n=== Mapping Summary ===")
        print(f"  Total entries: {len(result):,}")
        print(f"    - UniProt annotations: {annotation_count - no_annotation_count:,}")
        print(f"    - Miscleavage features: {miscleavage_count:,}")
        print(f"    - No annotation: {no_annotation_count:,}")
        print(f"  Processing time: {t1-t0:.2f} seconds")

    return result[[
        'Protein', 'PeptideID', 'pepStart', 'pepEnd', 
        'feature_category', 'feature', 'startpos', 'endpos', 
        'group', 'agent', 'note', 'score', 'dPF'
    ]].sort_values(['Protein', 'PeptideID', 'feature', 'startpos']).reset_index(drop=True)

peptide_uniprot_map = create_enhanced_peptide_uniprot_map_with_miscleavage(
    summary_data, uniprot_data, protein_dpf_analysis, 
    overlap_threshold=0.5, verbose=True
)
print()
print("Sample of Peptide-UniProt Mapping with Miscleavage Detection:")
peptide_uniprot_map.head(10)

# Filter for only PTM peptides (dPF == -1)
final_res = peptide_uniprot_map[
    (peptide_uniprot_map['dPF'] == -1)
]

# Number of Proteins with No Annotation
no_annot_proteins = final_res[final_res['feature'] == 'No Annotation']['Protein'].unique()
with_annotation_proteins = final_res[final_res['feature'] != 'No Annotation']['Protein'].unique()
print(f"Number of Proteins with No Annotation: {len(no_annot_proteins)}")
print(f"Number of Proteins with Annotation: {len(with_annotation_proteins)}")

# Create a donut chart
data = [len(no_annot_proteins), len(with_annotation_proteins)]
# Labels
labels = ["No Annotation", "With Annotation"]
# Colors
colors = ["#e5e5e5", "#a7c957"]

fig, ax = plt.subplots(figsize=(8, 6))
wedges, texts = ax.pie(
    data,
    labels=None,             
    startangle=140,
    colors=colors,
    wedgeprops=dict(width=0.60, edgecolor='w'),
    explode=[0.02, 0.0],  # separate the first slice slightly
    counterclock=False
)
# Add percent labels placed on wedges
bbox_props = dict(boxstyle="round,pad=0.3", fc="white", ec="none", alpha=0.8)
for w, p in zip(wedges, np.array(data) / sum(data)):
    angle = (w.theta2 + w.theta1) / 2.0
    x, y = w.r * 0.75 * np.cos(np.deg2rad(angle)), w.r * 0.75 * np.sin(np.deg2rad(angle))
    ax.text(x, y, f"{p:.1%}", ha='center', va='center', fontsize=15, bbox=bbox_props, fontweight='bold')
# Center annotation: total proteins
ax.text(0, 0, f"Total\n{sum(data)}", ha='center', va='center', fontsize=15, fontweight='bold')
# Legend with counts and percentages
legend_labels = [f"{lab}: {cnt:,} ({pct:.1%})" for lab, cnt, pct in zip(labels, data, np.array(data) / sum(data))]
ax.legend(
    wedges, legend_labels,
    title="Number of Proteins",
    loc="upper center",
    bbox_to_anchor=(0.5, -0.12),
    ncol=len(labels),
    fontsize=11,
    title_fontsize=12,
    frameon=False, 
    handletextpad=0.5,
    columnspacing=1.5,
)
ax.axis('equal')
plt.title('Distribution of Proteins with PTM Peptides\nAnnotation Availability', fontsize=16, fontweight='bold')
plt.tight_layout()
plots.finalize_plot( 
    fig, show=True, save=save_to_folder, 
    filename='proteins_withPTM_peptides_UniprotAnnotation_Donut',
    filepath=figure_path,
    formats=figure_formats,
    transparent=transparent_bg,
    dpi=figure_dpi
)

print()

# Create a visualization to show the Number of Proteins with UniProt Annotations
fig, ax = plt.subplots(figsize=(5, 5))
sub_data = final_res[final_res['feature'] != 'No Annotation']

# Create a bar plot
sub_data['feature'].value_counts().plot(kind='barh', color="#a7c957", ax=ax, edgecolor='black')
# Add counts to on top of the bars
for i, v in enumerate(sub_data['feature'].value_counts()):
    ax.text(
        # v, i, str(v),
        # Adjust the position of the text
        v + 2.5, i-0.1, str(v),
        color='black', 
        ha='left', va='center', 
        rotation=0, fontsize=12, fontweight='bold'
    )
ax.set_xlabel("Number of Annotations", fontsize=12)
ax.set_ylabel("Uniprot Feature Type", fontsize=12)
ax.set_title("Annotations Matched to Single PTMs", fontsize=14, fontweight='bold')
ax.grid("both", linestyle="--", linewidth=0.75, alpha=0.5, color="lightgrey")
plt.tight_layout()
plt.show()
plots.finalize_plot( 
    fig, show=True, save=save_to_folder, 
    filename='proteins_withPTM_peptides_UniprotAnnotation_Bar',
    filepath=figure_path,
    formats=figure_formats,
    transparent=transparent_bg,
    dpi=figure_dpi
)
print()
print("Modified Residues (MOD_RES) Annotations:")
print((sub_data.loc[sub_data['feature']=='MOD_RES', 'group'].str.split(';').str[0]).value_counts())
print()
print("Cleavage Annotations:")
print((sub_data.loc[sub_data['feature']=='CLEAVAGE', 'group'].str.split(';').str[0]).value_counts())
print()